{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ea857b1",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation(RAG) с использованием GigaChat на примере задачи \"разговор с книгой\"\n",
    "\n",
    "Подход RAG позволяет большим языковым моделям (LLM) отвечать на вопросы по документам, которы не помещаются в промпт.\n",
    "Ниже приведен пример того, как можно научить модель отвечать на вопросы, используя текст из книги \"Мастер и Маргарита\".\n",
    "\n",
    "Подробнее про RAG вы можете прочитать в [документации LangChain](https://python.langchain.com/docs/use_cases/question_answering/) и в курсе по промпт-инженирингу от Сбера (ссылка будет позже).\n",
    "\n",
    "В качестве примера мы рассмотрим текст романа Булгакова \"Мастер и Маргарита\" (главы 1 и 2).\n",
    "\n",
    "Вопрос будет - `Какой плащ был у Понтия Пилата?`. Ответ содержится во второй главе книги:\n",
    "`В белом плаще с кровавым подбоем, шаркающей кавалерийской походкой, ранним утром четырнадцатого числа весеннего месяца нисана в крытую колоннаду между двумя крыльями дворца ирода великого вышел прокуратор Иудеи Понтий Пилат.`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea31150",
   "metadata": {},
   "source": [
    "## Установка\n",
    "\n",
    "Для работы нам понадобится векторая база данных. Мы будем использовать Chroma."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7dc1ec5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install chromadb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7543fa",
   "metadata": {},
   "source": [
    "## Инициализация модели\n",
    "Теперь инициализируем модель GigaChat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "921afeaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models.gigachat import GigaChat\n",
    "\n",
    "llm = GigaChat(credentials=...)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84443c6f",
   "metadata": {},
   "source": [
    "Для проверки спросим у модели вопрос про цвет плаща без какого-либо контекста. Возможно, она и так будет давать ожидаемый ответ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "af0398d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Понтий Пилат, римский прокуратор Иудеи, известен своим красно-белым плащом. Этот плащ был характерной чертой его должности и символизировал его власть и статус. \\n\\nВ исторических источниках описывается'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.schema import HumanMessage\n",
    "\n",
    "question = \"Какой плащ был у Понтия Пилата?\"\n",
    "llm([HumanMessage(content=question)]).content[0:200]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b2e4e31",
   "metadata": {},
   "source": [
    "Видим, что модель не отвечает так, как нам хотелось бы, поэтому применим RAG-подход.\n",
    "\n",
    "## Подготовка документа\n",
    "\n",
    "Для работы с документом нам нужно разделить его на части. Для этого используем `TextLoader` для загрузки книги и `RecursiveCharacterTextSplitter`, чтобы разделить текст на приблизительно равные куски в ~1000 символов с перекрытием в ~200 символов. Этот тип сплиттера сам выбирает каким способом следует оптимально разделять документ (по абзацам, по предложениям и т.д.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f8cf5765",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total documents: 91\n"
     ]
    }
   ],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.text_splitter import (\n",
    "    RecursiveCharacterTextSplitter,\n",
    ")\n",
    "\n",
    "loader = TextLoader(\"../мастер_и_маргарита.txt\")\n",
    "documents = loader.load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=1000,\n",
    "    chunk_overlap=200,\n",
    ")\n",
    "documents = text_splitter.split_documents(documents)\n",
    "print(f\"Total documents: {len(documents)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131d5059",
   "metadata": {},
   "source": [
    "После нарезки мы получили 91 документ частями книги.\n",
    "\n",
    "## Создание базы данных эмбеддингов\n",
    "\n",
    "Эмбеддинг это векторное представление текста, которое может быть использовано для определения смысловой близости текстов. Векторная база данных хранит тексты и соответствующие им эмбеддинги, а также умеет выполнять поиск по ним. Для работы с базой данных мы создаем объект GigaChatEmbeddings и передаем его в базу данных Chroma.\n",
    "\n",
    "> Сервис эмбеддингов на данный момент находится на этапе предварительного тестирования и может быть не доступен всем пользователям или не стабилен.\n",
    "\n",
    "> Параметры `one_by_one_mode` и `_debug_delay` используются для работы в режиме отладки и снижения нагрузки на сервер работы с эммбеддингами."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "fdce8923",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to send telemetry event ClientStartEvent: module 'chromadb' has no attribute '__version__'\n",
      "Failed to send telemetry event ClientCreateCollectionEvent: module 'chromadb' has no attribute '__version__'\n"
     ]
    }
   ],
   "source": [
    "from chromadb.config import Settings\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain_community.embeddings import GigaChatEmbeddings\n",
    "\n",
    "embeddings = GigaChatEmbeddings(credentials=\"...\")\n",
    "\n",
    "db = Chroma.from_documents(\n",
    "    documents,\n",
    "    embeddings,\n",
    "    client_settings=Settings(anonymized_telemetry=False),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29137915",
   "metadata": {},
   "source": [
    "## Поиск по базе данных\n",
    "\n",
    "Теперь можно обратиться к базе данных и попросить найти документы, которые с наибольшей вероятностью содержат ответ на наш вопрос.\n",
    "\n",
    "По-умолчанию база данных возвращает 4 наиболее релевантных документа. Этот параметр можно изменить в зависимости от решаемой задачи и типа документов.\n",
    "\n",
    "Видно, что первый же документ содержит внутри себя часть книги с ответом на наш вопрос."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b0c55e98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs = db.similarity_search(question, k=4)\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "32b43339",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "... акцент почему-то пропал: – Все просто: в белом плаще...\\n\\n\\n\\nГлава 2\\n\\nПонтий Пилат\\n\\nВ белом плаще с кровавым подбоем, шаркающей кавалерийской походкой, ранним утром четырнадц ...\n"
     ]
    }
   ],
   "source": [
    "print(f\"... {str(docs[0])[620:800]} ...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a063c7",
   "metadata": {},
   "source": [
    "## QnA цепочка\n",
    "\n",
    "Теперь мы создадим цепочку QnA, которая специально предназначена для ответов на вопросы по документам. В качестве аргументов есть передается языковая модель и ретривер (база данных)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "24671e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "qa_chain = RetrievalQA.from_chain_type(llm, retriever=db.as_retriever())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "592151a2",
   "metadata": {},
   "source": [
    "Наконец можно задать вопрос нашей цепочке и получить правильный ответ!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f3b2f453",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'Какой плащ носил Понтий Пилат?',\n",
       " 'result': 'Понтий Пилат носил белый плащ с кровавым подбоем.'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_chain({\"query\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa10ec4",
   "metadata": {},
   "source": [
    "Несколько дополнительных вопросов для проверки работоспособности:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "8f619e09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'Какая трость была у Воланда?',\n",
       " 'result': 'Трость, которую держал Воланд, имела черный набалдашник в виде головы пуделя.'}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_chain({\"query\": \"Какая трость была у Воланда?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a4a54b82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'В чем главная проблема человека?',\n",
       " 'result': 'Главная проблема человека, согласно данному контексту, заключается в том, что он не может управлять своей жизнью и всем распорядком на земле без точного плана на некоторый срок.'}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa_chain({\"query\": \"В чем главная проблема человека?\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
